{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "456f6c6b",
   "metadata": {},
   "source": [
    "# 에이전트 대화 시뮬레이션(고객 응대 시나리오)\n",
    "\n",
    "챗봇을 구축할 때, 예를 들어 고객 지원 어시스턴트와 같은 경우, 챗봇의 성능을 제대로 평가하는 것이 어려울 수 있습니다. 코드 변경마다 집중적으로 수동으로 상호 작용하는 것은 시간이 많이 소요됩니다.\n",
    "\n",
    "평가 과정을 더 쉽고 재현 가능하게 만드는 한 가지 방법은 **사용자 상호 작용을 시뮬레이션하는 것** 입니다.\n",
    "\n",
    "LangGraph를 사용하면 이를 설정하는 것이 쉽습니다. \n",
    "\n",
    "아래는 대화를 시뮬레이션하기 위해 \"가상 사용자(Simulated User)\"를 생성하는 방법의 예시입니다.\n",
    "\n",
    "![agent-simulations.png](assets/agent-simulations.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b052df4d",
   "metadata": {},
   "source": [
    "## 환경 설정"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12ddac30",
   "metadata": {},
   "outputs": [],
   "source": [
    "# API 키를 환경변수로 관리하기 위한 설정 파일\n",
    "from dotenv import load_dotenv\n",
    "\n",
    "# API 키 정보 로드\n",
    "load_dotenv()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a69d4c5e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# LangSmith 추적을 설정합니다. https://smith.langchain.com\n",
    "# !pip install -qU langchain-teddynote\n",
    "from langchain_teddynote import logging\n",
    "\n",
    "# 프로젝트 이름을 입력합니다.\n",
    "logging.langsmith(\"CH17-LangGraph-Use-Cases\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "928f422d",
   "metadata": {},
   "source": [
    "## 상태(State) 정의"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "70538645",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langgraph.graph.message import add_messages\n",
    "from typing import Annotated\n",
    "from typing_extensions import TypedDict\n",
    "\n",
    "\n",
    "# State 정의\n",
    "class State(TypedDict):\n",
    "    messages: Annotated[list, add_messages]  # 사용자 - 상담사 간의 대화 메시지"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "11f305bf",
   "metadata": {},
   "source": [
    "## 상담사, 고객 역할 정의\n",
    "\n",
    "### 상담사 역할 정의\n",
    "\n",
    "시뮬레이션에서 **상담사** 역할을 하는 챗봇을 정의합니다.\n",
    "\n",
    "**참고**\n",
    "\n",
    "- `call_chatbot` 내의 구현은 설정 가능하며, 내부에서 사용한 모델을 Agent 로 변경하는 것도 가능합니다.\n",
    "- `call_chatbot` 은 사용자로부터 메시지를 입력으로 받아, 고객을 상담하는 역할을 부여하겠습니다. \n",
    "\n",
    "*고객 지원 시나리오에서의 대화 응답 생성에 활용될 수 있습니다.*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "ce12ed49",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "from langchain_teddynote.models import LLMs, get_model_name\n",
    "from langchain_openai import ChatOpenAI\n",
    "from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder\n",
    "from langchain_core.messages import HumanMessage, AIMessage, BaseMessage\n",
    "from langchain_core.output_parsers import StrOutputParser\n",
    "\n",
    "# 모델 이름 설정\n",
    "MODEL_NAME = get_model_name(LLMs.GPT4)\n",
    "\n",
    "\n",
    "def call_chatbot(messages: List[BaseMessage]) -> dict:\n",
    "    # LangChain ChatOpenAI 모델을 Agent 로 변경할 수 있습니다.\n",
    "    prompt = ChatPromptTemplate.from_messages(\n",
    "        [\n",
    "            (\n",
    "                \"system\",\n",
    "                \"You are a customer support agent for an airline. Answer in Korean.\",\n",
    "            ),\n",
    "            MessagesPlaceholder(variable_name=\"messages\"),\n",
    "        ]\n",
    "    )\n",
    "    model = ChatOpenAI(model=MODEL_NAME, temperature=0.6)\n",
    "    chain = prompt | model | StrOutputParser()\n",
    "    return chain.invoke({\"messages\": messages})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dcc866c9",
   "metadata": {},
   "source": [
    "`call_chatbot` 은 사용자의 입력을 받아 챗봇의 응답을 처리합니다. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "afd97af6",
   "metadata": {},
   "outputs": [],
   "source": [
    "call_chatbot([(\"user\", \"안녕하세요?\")])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "99b16074",
   "metadata": {},
   "source": [
    "### 고객 역할(Simulated User) 정의\n",
    "\n",
    "이제 시뮬레이션된 고객의 역할을 정의합니다. 고객 지원 시나리오에서의 대화를 시뮬레이션합니다. \n",
    "\n",
    "시스템 프롬프트는 고객과 고객 지원 담당자 간의 상호작용을 설정하며, 사용자 지시사항을 통해 시나리오의 세부 사항을 제공합니다. \n",
    "\n",
    "이 구성은 특정 사용자 요구(예: 환불 요청)에 대한 모델의 반응을 시뮬레이션하는 데 사용됩니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "142d4247",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder\n",
    "from langchain_openai import ChatOpenAI\n",
    "\n",
    "\n",
    "def create_scenario(name: str, instructions: str):\n",
    "    # 시스템 프롬프트를 정의: 필요에 따라 변경\n",
    "    system_prompt_template = \"\"\"You are a customer of an airline company. \\\n",
    "You are interacting with a user who is a customer support person. \\\n",
    "\n",
    "Your name is {name}.\n",
    "\n",
    "# Instructions:\n",
    "{instructions}\n",
    "\n",
    "[IMPORTANT] \n",
    "- When you are finished with the conversation, respond with a single word 'FINISHED'\n",
    "- You must speak in Korean.\"\"\"\n",
    "\n",
    "    # 대화 메시지와 시스템 프롬프트를 결합하여 채팅 프롬프트 템플릿을 생성합니다.\n",
    "    prompt = ChatPromptTemplate.from_messages(\n",
    "        [\n",
    "            (\"system\", system_prompt_template),\n",
    "            MessagesPlaceholder(variable_name=\"messages\"),\n",
    "        ]\n",
    "    )\n",
    "\n",
    "    # 특정 사용자 이름과 지시사항을 사용하여 프롬프트를 부분적으로 채웁니다.\n",
    "    prompt = prompt.partial(name=name, instructions=instructions)\n",
    "    return prompt"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "212913b9",
   "metadata": {},
   "source": [
    "가상의 시나리오를 생성합니다. 이 가상의 시나리오는 고객의 입장에서의 시나리오입니다.\n",
    "\n",
    "여기서는 환불을 요청하는 시나리오를 정의합니다.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "691500b9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 사용자 지시사항을 정의합니다.\n",
    "instructions = \"\"\"You are trying to get a refund for the trip you took to Jeju Island. \\\n",
    "You want them to give you ALL the money back. This trip happened last year.\"\"\"\n",
    "\n",
    "# 사용자 이름을 정의합니다.\n",
    "name = \"Teddy\"\n",
    "\n",
    "create_scenario(name, instructions).pretty_print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "6fcfb502",
   "metadata": {},
   "outputs": [],
   "source": [
    "# OpenAI 챗봇 모델을 초기화합니다.\n",
    "model = ChatOpenAI(model=MODEL_NAME, temperature=0.6)\n",
    "\n",
    "# 시뮬레이션된 사용자 대화를 생성합니다.\n",
    "simulated_user = create_scenario(name, instructions) | model | StrOutputParser()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9111fd1c",
   "metadata": {},
   "source": [
    "생성된 `simulated_user` 를 호출하여 시뮬레이션된 사용자에게 메시지를 전달합니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e5349086",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.messages import HumanMessage\n",
    "\n",
    "# 시뮬레이션된 사용자에게 메시지를 전달(상담사 -> 고객)\n",
    "messages = [HumanMessage(content=\"안녕하세요? 어떻게 도와 드릴까요?\")]\n",
    "simulated_user.invoke({\"messages\": messages})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d5a96c1d",
   "metadata": {},
   "source": [
    "## 에이전트 시뮬레이션 정의하기\n",
    "\n",
    "아래의 코드는 시뮬레이션을 실행하기 위한 LangGraph 워크플로우를 생성합니다. \n",
    "\n",
    "주요 구성 요소는 다음과 같습니다:\n",
    "\n",
    "1. 시뮬레이션된 사용자와 챗봇을 위한 두 개의 노드입니다.\n",
    "2. 조건부 정지 기준을 가진 그래프 자체입니다."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5abdc636",
   "metadata": {},
   "source": [
    "### 노드 정의\n",
    "\n",
    "먼저, 그래프에서 노드를 정의합니다. 이들은 메시지 목록을 입력으로 받아 상태에 추가할 메시지 목록을 반환해야 합니다.\n",
    "이것들은 위에 있는 챗봇과 시뮬레이션된 사용자를 둘러싼 것 래퍼들입니다.\n",
    "\n",
    "**참고:** 여기서 까다로운 점은 어떤 메시지가 어떤 것인지 구분하는 것입니다. \n",
    "\n",
    "챗봇과 시뮬레이션된 사용자 모두 LLMs이기 때문에, 둘 다 AI 메시지로 응답할 것입니다. 우리의 상태는 인간과 AI 메시지가 번갈아 가며 나열된 목록이 될 것입니다. 이는 노드 중 하나에서 AI와 인간 역할을 바꾸는 논리가 필요함을 의미합니다. \n",
    "\n",
    "이 예제에서는, HumanMessages가 시뮬레이션된 사용자로부터 온 메시지라고 가정할 것입니다. 이는 시뮬레이션된 사용자 노드에 AI와 Human 메시지를 교환하는 논리가 필요함을 의미합니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "41154718",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.messages import AIMessage\n",
    "\n",
    "\n",
    "# 상담사 역할\n",
    "def ai_assistant_node(messages):\n",
    "    # 상담사 응답 호출\n",
    "    ai_response = call_chatbot(messages)\n",
    "\n",
    "    # AI 상담사의 응답을 반환\n",
    "    return {\"messages\": [(\"assistant\", ai_response)]}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3a0263be",
   "metadata": {},
   "source": [
    "상담사 역할의 노드를 호출합니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d3b77124",
   "metadata": {},
   "outputs": [],
   "source": [
    "ai_assistant_node(\n",
    "    [\n",
    "        (\"user\", \"안녕하세요?\"),\n",
    "        (\"assistant\", \"안녕하세요! 어떻게 도와드릴까요?\"),\n",
    "        (\"user\", \"환불 어떻게 하나요?\"),\n",
    "    ]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b07573b6",
   "metadata": {},
   "source": [
    "다음으로, 우리의 시뮬레이션된 사용자를 위한 노드를 정의해 보겠습니다.\n",
    "\n",
    "**참고**\n",
    "\n",
    "- 이 과정에서는 메시지의 역할을 교체하는 작은 로직이 포함될 것입니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "c2654c6f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def _swap_roles(messages):\n",
    "    # 메시지의 역할을 교환: 시뮬레이션 사용자 단계에서 메시지 타입을 AI -> Human, Human -> AI 로 교환합니다.\n",
    "    new_messages = []\n",
    "    for m in messages:\n",
    "        if isinstance(m, AIMessage):\n",
    "            # AIMessage 인 경우, HumanMessage 로 변환합니다.\n",
    "            new_messages.append(HumanMessage(content=m.content))\n",
    "        else:\n",
    "            # HumanMessage 인 경우, AIMessage 로 변환합니다.\n",
    "            new_messages.append(AIMessage(content=m.content))\n",
    "    return new_messages\n",
    "\n",
    "\n",
    "# 상담사 역할(AI Assistant) 노드 정의\n",
    "def ai_assistant_node(state: State):\n",
    "    # 상담사 응답 호출\n",
    "    ai_response = call_chatbot(state[\"messages\"])\n",
    "\n",
    "    # AI 상담사의 응답을 반환\n",
    "    return {\"messages\": [(\"assistant\", ai_response)]}\n",
    "\n",
    "\n",
    "# 시뮬레이션된 사용자(Simulated User) 노드 정의\n",
    "def simulated_user_node(state: State):\n",
    "    # 메시지 타입을 교환: AI -> Human, Human -> AI\n",
    "    new_messages = _swap_roles(state[\"messages\"])\n",
    "\n",
    "    # 시뮬레이션된 사용자를 호출\n",
    "    response = simulated_user.invoke({\"messages\": new_messages})\n",
    "    return {\"messages\": [(\"user\", response)]}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "01ec6e36",
   "metadata": {},
   "source": [
    "### 엣지 정의\n",
    "\n",
    "이제 엣지에 대한 로직을 정의할 필요가 있습니다. 주된 로직은 시뮬레이션된 사용자가 작업을 마친 후 발생하며, 두 가지 결과 중 하나로 이어져야 합니다:\n",
    "\n",
    "- 고객 지원 봇을 호출하여 계속 진행(\"continue\")\n",
    "- 대화를 마치고 종료(\"end\")\n",
    "\n",
    "그렇다면 대화가 종료되는 로직은 무엇일까요? 우리는 이를 인간 챗봇이 `FINISHED`로 응답하거나(시스템 프롬프트 참조) 대화가 6개 메시지를 초과하는 경우로 정의할 것입니다 (이는 이 예제를 짧게 유지하기 위한 임의의 숫자입니다)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "10836dc2",
   "metadata": {},
   "source": [
    "`should_continue` 함수는 메시지 리스트를 인자로 받아, 리스트의 길이가 6을 초과하거나 마지막 메시지의 내용이 'FINISHED'일 경우 'end'를 반환합니다. \n",
    "\n",
    "그렇지 않으면 'continue'를 반환하여 처리를 계속하도록 합니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "988b59ab",
   "metadata": {},
   "outputs": [],
   "source": [
    "def should_continue(state: State):\n",
    "    # 메시지 리스트의 길이가 6보다 크면 'end'를 반환합니다.\n",
    "    if len(state[\"messages\"]) > 6:\n",
    "        return \"end\"\n",
    "    # 마지막 메시지의 내용이 'FINISHED'라면 'end'를 반환합니다.\n",
    "    elif state[\"messages\"][-1].content == \"FINISHED\":\n",
    "        return \"end\"\n",
    "    # 위의 조건에 해당하지 않으면 'continue'를 반환합니다.\n",
    "    else:\n",
    "        return \"continue\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "48a99152",
   "metadata": {},
   "source": [
    "## 그래프 정의\n",
    "\n",
    "이제 시뮬레이션을 설정하는 그래프를 정의합니다.\n",
    "\n",
    "`MessageGraph` 클래스는 챗봇과 시뮬레이션된 사용자 간의 상호작용을 구성하고 시뮬레이션하는 데 사용됩니다. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "562a67cd",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langgraph.graph import END, StateGraph\n",
    "\n",
    "# StateGraph 인스턴스 생성\n",
    "graph_builder = StateGraph(State)\n",
    "\n",
    "# 노드 정의\n",
    "graph_builder.add_node(\"simulated_user\", simulated_user_node)\n",
    "graph_builder.add_node(\"ai_assistant\", ai_assistant_node)\n",
    "\n",
    "# 엣지 정의 (챗봇 -> 시뮬레이션된 사용자)\n",
    "graph_builder.add_edge(\"ai_assistant\", \"simulated_user\")\n",
    "\n",
    "# 조건부 엣지 정의\n",
    "graph_builder.add_conditional_edges(\n",
    "    \"simulated_user\",\n",
    "    should_continue,\n",
    "    {\n",
    "        \"end\": END,  # 종료 조건이 충족되면 시뮬레이션을 중단\n",
    "        \"continue\": \"ai_assistant\",  # 종료 조건이 충족되지 않으면 상담사 역할 노드로 메시지를 전달\n",
    "    },\n",
    ")\n",
    "\n",
    "# 시작점 설정\n",
    "graph_builder.set_entry_point(\"ai_assistant\")\n",
    "\n",
    "# 그래프 컴파일\n",
    "simulation = graph_builder.compile()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e0e1f8b8",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_teddynote.graphs import visualize_graph\n",
    "\n",
    "visualize_graph(simulation)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "854607cb",
   "metadata": {},
   "source": [
    "## 시뮬레이션 시작\n",
    "\n",
    "이제 우리의 챗봇을 평가할 수 있습니다! 빈 메시지로 호출할 수 있습니다(이것은 챗봇이 초기 대화를 시작하게 하는 것을 시뮬레이션합니다)\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "777dd349",
   "metadata": {},
   "source": [
    "시뮬레이션에서 스트리밍되는 데이터 청크를 순회하며, 최종 종료 청크(`END`)를 제외한 모든 이벤트를 출력합니다. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dccd4d54",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.runnables import RunnableConfig\n",
    "from langchain_teddynote.messages import stream_graph, random_uuid\n",
    "\n",
    "\n",
    "# config 설정(재귀 최대 횟수, thread_id)\n",
    "config = RunnableConfig(recursion_limit=10, configurable={\"thread_id\": random_uuid()})\n",
    "\n",
    "# 입력 메시지 설정\n",
    "inputs = {\n",
    "    \"messages\": [HumanMessage(content=\"안녕하세요? 저 지금 좀 화가 많이 났습니다^^\")]\n",
    "}\n",
    "\n",
    "# 그래프 스트리밍\n",
    "stream_graph(simulation, inputs, config, node_names=[\"simulated_user\", \"ai_assistant\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c277fc53",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "langchain-kr-lwwSZlnu-py3.11",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
